require 'zlib'
require 'pathname'

class MnistSuite
  include Numo

  def training_data
    @training_set ||= training_images.zip(training_labels)
  end


  def validation_data
    @validation_set ||= validation_images.zip(validation_labels)
  end


  def test_data
    @test_set ||= test_images.zip(test_labels)
  end

  private

  def training_images
    read_images(path_for 'train-images-idx3-ubyte.gz')[0 ... 50000]
  end

  def training_labels
    read_labels(path_for 'train-labels-idx1-ubyte.gz')[0 ... 50000]
  end

  def validation_images
    read_images(path_for 'train-images-idx3-ubyte.gz')[50000 .. -1]
  end

  def validation_labels
    read_labels(path_for 'train-labels-idx1-ubyte.gz')[50000 ... -1]
  end

  def test_images
    read_images(path_for 't10k-images-idx3-ubyte.gz')
  end

  def test_labels
    read_labels(path_for 't10k-labels-idx1-ubyte.gz')
  end

  def read_images(path)
    Zlib::GzipReader.open(path) do |gz|
      content = gz.read.unpack('l>4C*')
      pixels_per_image = content[2] * content[3]
      content[4 .. -1].map!{|n| n / 255.0}.each_slice(pixels_per_image).map{|img| DFloat[img].transpose}
    end
  end

  def read_labels(path)
    Zlib::GzipReader.open(path) do |gz|
      gz.read.unpack('l>2C*')[2 .. -1].map! do |n|
        DFloat[Array.new(10) {|i| i == n ? 1 : 0}].transpose
      end
    end
  end

  def path_for(filename)
    Pathname(File.expand_path('../../data', __FILE__)).join(filename)
  end
end
